#include <proc.h>
#include <slab.h>
#include <string.h>
#include <sync.h>
#include <pmm.h>
#include <error.h>
#include <sched.h>
#include <elf.h>
#include <vmm.h>
#include <trap.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <um.h>
#include <stub.h>
#include <host_syscall.h>
#include <swap.h>

void
switch_to(struct context *from, struct context *to) 
{
	if (setjmp (from->switch_buf) == 0)
		longjmp (to->switch_buf, 1);
}

// alloc_proc - create a proc struct and init fields
struct proc_struct *
alloc_proc(void) {
    struct proc_struct *proc = kmalloc(sizeof(struct proc_struct));
    if (proc != NULL) {
        proc->state = PROC_UNINIT;
        proc->pid = -1;
        proc->runs = 0;
        proc->kstack = 0;
        proc->need_resched = 0;
        proc->parent = NULL;
        proc->mm = NULL;
        memset(proc->name, 0, PROC_NAME_LEN);
        proc->wait_state = 0;
        proc->cptr = proc->optr = proc->yptr = NULL;
        list_init(&(proc->thread_group));
		
        memset(&(proc->context), 0, sizeof(struct context));
        proc->tf = NULL;
        proc->flags = 0;
		
        proc->rq = NULL;
        list_init(&(proc->run_link));
        proc->time_slice = 0;

		proc->arch.host = NULL;
		proc->arch.forking = 0;
        proc->sem_queue = NULL;
        event_box_init(&(proc->event_box));
    }
    return proc;
}

// kernel_thread - create a kernel thread using "fn" function
// NOTE: the contents of temp trapframe tf will be copied to 
//       proc->tf in do_fork-->copy_thread function
int
kernel_thread(int (*fn)(void *), void *arg, uint32_t clone_flags) {
	struct trapframe tf;
	memset (&tf, 0, sizeof (struct trapframe));
	tf.fn = fn;
	tf.arg = arg;
	return do_fork (clone_flags | CLONE_VM, 0, &tf);
}

/**
 * de_thread_arch_hook - called everytime @de_thread is called to perform arch-related operations.
 */
void
de_thread_arch_hook (struct proc_struct *proc) {
	if (proc->arch.host != NULL) {
		assert (proc->arch.host->nr_threads > 0);
		proc->arch.host->nr_threads --;
		if (proc->arch.host->nr_threads == 0) {
			free_page (kva2page (current->arch.host->stub_stack));
			kill_ptraced_process (current->arch.host->host_pid, 1);
			kfree (current->arch.host);
		}
	}
}

// copy_thread - setup the trapframe on the  process's kernel stack top and
//             - setup the kernel entry point and stack of process
int
copy_thread(uint32_t clone_flags, struct proc_struct *proc,
			uintptr_t user_stack, struct trapframe *tf) {
	int pid;
	void *stub_stack;
	
	if (RUN_AS_USER) {
		if (clone_flags & CLONE_VM) {
			/* Use current host process as its container */
			proc->arch.host = current->arch.host;
			proc->arch.host->nr_threads ++;
		} else {
			/* Create a new child process */
			proc->arch.host = kmalloc (sizeof (struct host_container));
			if (proc->arch.host == NULL)
				goto exit;
			
			stub_stack = boot_alloc_page ();
			if (stub_stack == NULL)
				goto exit_free_host;
			pid = start_userspace (stub_stack);
		
			if (pid < 0)
				goto exit_free_stub;
			
			proc->arch.host->host_pid = pid;
			proc->arch.host->nr_threads = 1;
			proc->arch.host->stub_stack = (struct stub_stack *)stub_stack;
			/* unmap kernel area. */
			if (host_munmap (proc, (void*)KERNBASE, KERNTOP-KERNBASE) < 0)
				panic ("unmap kernel area failed \n");
		}
		memcpy (&(proc->arch.regs.regs), &(current->arch.regs.regs), sizeof (struct user_regs_struct));
		proc->arch.regs.regs.eax = 0; /* make the child return 0 for the syscall */
		proc->arch.regs.regs.esp = user_stack;
		
		/* The current thread will run in 'userspace' */
		proc->context.switch_buf->__ebx = (uint32_t)userspace;
		proc->context.switch_buf->__ebp = (uint32_t)&(proc->arch.regs);
	} else {
		proc->context.switch_buf->__ebx = (uint32_t)(tf->fn);
		proc->context.switch_buf->__ebp = (uint32_t)(tf->arg);
	}
	proc->context.switch_buf->__eip = (uint32_t)kernel_thread_entry;
	proc->context.switch_buf->__esp = proc->kstack + KSTACKSIZE;
	return 0;

exit_free_stub:
	free_page (kva2page (stub_stack));
exit_free_host:
	kfree (proc->arch.host);
exit:
	return -1;
}

void
init_new_context (struct proc_struct *proc, struct elfhdr *elf) {
	proc->arch.regs.regs.esp = USTACKTOP;
    proc->arch.regs.regs.eip = elf->e_entry;
	proc->arch.regs.regs.xcs = 0x73;
	proc->arch.regs.regs.xds = proc->arch.regs.regs.xes =
		proc->arch.regs.regs.xfs = proc->arch.regs.regs.xgs =
		proc->arch.regs.regs.xss = 0x7b;
}

int
do_execve_arch_hook ()
{
	current->arch.host = kmalloc (sizeof (struct host_container));
	if (current->arch.host == NULL)
		return -1;

	void* stub_stack = boot_alloc_page ();
	if (stub_stack == NULL)
		goto free_host;
	int ret = start_userspace (stub_stack);
	if (ret <= 0)
		goto free_stub_stack;

	current->arch.host->stub_stack = stub_stack;
	current->arch.host->host_pid = ret;
	current->arch.host->nr_threads = 1;
	tlb_invalidate_user ();
	
	/* unmap kernel area */
	if (host_munmap (current, (void*)KERNBASE, KERNTOP-KERNBASE) < 0)
		panic ("unmap kernel area failed\n");
	/* erase kernel maps in the page table */
	int valid_size = KERNBASE / PTSIZE * sizeof(pte_t);
	memset ((void*)((int)(current->mm->pgdir) + valid_size),
			0, PGSIZE - valid_size);
	userspace (&(current->arch.regs));
	/* should never comes here */
	
free_stub_stack:
	free_page (kva2page (stub_stack));
free_host:
	kfree (current->arch.host);
	current->arch.host = NULL;
	return -1;
}

// kernel_execve - exec a user program called by user_main kernel_thread
// different from x86's
int
kernel_execve(const char *name, unsigned char *binary, size_t size) {
    int ret, len = strlen(name);
    ret = do_execve (name, len, binary, size);
	return ret;
}
